(*
 * OWL - OCaml Scientific and Engineering Computing
 * Copyright (c) 2016-2020 Liang Wang <liang.wang@cl.cam.ac.uk>
 *)

type ('a, 'b) kind = ('a, 'b) Bigarray.kind

type ('a, 'b) t =
  { mutable m : int
  ; (* number of rows *)
    mutable n : int
  ; (* number of columns *)
    mutable k : ('a, 'b) kind
  ; (* type of sparse matrices *)
    mutable d : (int * int, 'a) Hashtbl.t (* hashtbl for storing data *)
  }

let zeros ?(density = 0.30) k m n =
  { m
  ; n
  ; k
  ; d =
      (let c = int_of_float (float_of_int (m * n) *. density) in
       Hashtbl.create c)
  }


let row_num x = x.m

let col_num x = x.n

let shape x = x.m, x.n

let numel x = x.m * x.n

let prune x r = Hashtbl.filter_map_inplace (fun _ v -> if v = r then None else Some v) x.d

let nnz x =
  let _ = prune x (Owl_const.zero x.k) in
  Hashtbl.((stats x.d).num_bindings)


let density x = float_of_int (nnz x) /. float_of_int (numel x)

let kind x = x.k

let _check_boundary i j m n =
  if i < 0 || i >= m || j < 0 || j >= n then failwith "error: index beyond the boundary"


let set x i j a =
  _check_boundary i j x.m x.n;
  let _a0 = Owl_const.zero x.k in
  match Hashtbl.mem x.d (i, j) with
  | true  -> if a <> _a0 then Hashtbl.replace x.d (i, j) a else Hashtbl.remove x.d (i, j)
  | false -> if a <> _a0 then Hashtbl.add x.d (i, j) a


let get x i j =
  _check_boundary i j x.m x.n;
  match Hashtbl.mem x.d (i, j) with
  | true  -> Hashtbl.find x.d (i, j)
  | false -> Owl_const.zero x.k


let reset x = Hashtbl.reset x.d

let copy x = { m = x.m; n = x.n; k = x.k; d = Hashtbl.copy x.d }

let iteri_nz f x = Hashtbl.iter (fun (i, j) v -> f i j v) x.d

let save x f = Owl_io.marshal_to_file x f

let load _k f = Owl_io.marshal_from_file f
